# -*- codeing:utf-8 -*-
"""
@author: jiaoyulong
@datetime: 2025/3/27 16:16
@Blog: 读取csv文件数据保存至本地数据库
"""
# 以下为准备工作
# 1.安装rarfile库
#   pip install rarfile
# 2.下载UnRAR工具
#   2.1访问RARLab官网下载UnRAR：https://www.rarlab.com/rar_add.htm
#   2.2选择UnRAR for Windows，解压下载的压缩包，找到unrar.exe
# 3.配置UnRAR路径(在python代码中指定路径)
#   import rarfile
#   rarfile.UNRAR_TOOL=C:\Path\To\unrar.exe

import shutil
from typing import Dict, Union, List

import pandas as pd
import os
import dolphindb as ddb
import rarfile
from sqlalchemy import create_engine   # SQLAlchemy==1.4.23


def connect_mysql(username: str, password: str, host: str, port: int, database: str):
    """
    创建连接
    :param username: 用户名
    :param password: 密码
    :param host: IP地址
    :param port: 端口
    :param database: 数据库名
    :return:
    """
    # 创建MySQL数据库连接字符串
    # engine = f'mysql+pymysql://{username}:{password}@{host}:{port}/{database}?charset=utf8mb4'
    engine = create_engine('mysql+pymysql://%s:%s@%s:%d/%s?charset=utf8mb4'
                           % (username, password, host, port, database))

    return engine


def remove_directory(directory):
    """
    删除文件夹及文件夹中的所有子文件
    :param directory: 指定目录路径
    """
    shutil.rmtree(directory)


# 保存至mysql数据库
def to_save_mysql_day(engine, df):
    table = 'future_day_data'

    # df.to_sql(name=table, con=engine, index=False, if_exists='replace')   # 如果表已存在，就替换
    df.to_sql(name=table, con=engine, index=False, if_exists='append')  # 添加现有表


def to_save_mysql_bar(engine, df):
    table = 'future_bar_data'

    # df.to_sql(name=table, con=engine, index=False, if_exists='replace')   # 如果表已存在，就替换
    df.to_sql(name=table, con=engine, index=False, if_exists='append')  # 添加现有表


def to_save_mysql_tick(engine, df):
    table = 'future_tick_data'

    # df.to_sql(name=table, con=engine, index=False, if_exists='replace')   # 如果表已存在，就替换
    df.to_sql(name=table, con=engine, index=False, if_exists='append')  # 添加现有表


# 保存至dolphindb数据库
def to_save_dolphindb_day(s, df):
    db_path = 'dfs://futures_day_level'
    table_name = 'futures_day'

    if not s.existsTable(db_path, table_name):
        print(f'不存在：{db_path}-->{table_name}')
        return

    df['TradingDay'] = pd.to_datetime(df['TradingDay'])

    # 更新数据
    upsert = ddb.tableUpsert(db_path, table_name, s,
                             keyColNames=['InstrumentID', 'TradingDay'])  # keyColNames设置去重列
    upsert.upsert(df)


def to_save_dolphindb_bar(s, df):
    db_path = 'dfs://futures_bar_level'
    table_name = 'futures_bar'

    if not s.existsTable(db_path, table_name):
        print(f'不存在：{db_path}-->{table_name}')
        return

    df['TradingDay'] = pd.to_datetime(df['TradingDay'])
    df['DateTime'] = pd.to_datetime(df['DateTime'])

    # 更新数据
    upsert = ddb.tableUpsert(db_path, table_name, s,
                             keyColNames=['InstrumentID', 'DateTime'])  # keyColNames设置去重列
    upsert.upsert(df)


def to_save_dolphindb_tick(s, df):
    db_path = 'dfs://futures_tick_level'
    table_name = 'futures_tick'

    if not s.existsTable(db_path, table_name):
        print(f'不存在：{db_path}-->{table_name}')
        return

    df['TradingDay'] = pd.to_datetime(df['TradingDay'])
    df['DateTime'] = pd.to_datetime(df['DateTime'])

    df['AskPrice'] = df[['AskPrice1', 'AskPrice2', 'AskPrice3', 'AskPrice4', 'AskPrice5']].values.tolist()
    df['AskVolume'] = df[['AskVolume1', 'AskVolume2', 'AskVolume3', 'AskVolume4', 'AskVolume5']].values.tolist()
    df['BidPrice'] = df[['BidPrice1', 'BidPrice2', 'BidPrice3', 'BidPrice4', 'BidPrice5']].values.tolist()
    df['BidVolume'] = df[['BidVolume1', 'BidVolume2', 'BidVolume3', 'BidVolume4', 'BidVolume5']].values.tolist()

    df = df[['InstrumentID', 'ExchangeID', 'TradingDay', 'DateTime', 'OpenPrice', 'HighestPrice', 'LowestPrice',
             'LastPrice', 'PreClosePrice', 'Turnover', 'Volume', 'UpperLimitPrice', 'LowerLimitPrice', 'OpenInterest',
             'AskPrice', 'AskVolume', 'BidPrice', 'BidVolume', 'PreSettlementPrice']]

    if not s.existsTable(db_path, table_name):
        print(f'不存在：{db_path}-->{table_name}')
        return

    # 更新数据
    upsert = ddb.tableUpsert(db_path, table_name, s,
                             keyColNames=['InstrumentID', 'DateTime'])  # keyColNames设置去重列
    upsert.upsert(df)


def get_all_files(folder_path):
    '''
    获取所有的rar文件
    :param folder_path: 路径
    :return:
    '''
    file_list = []
    for root, dirs, files in os.walk(folder_path):
        for file_name in files:
            file_path = os.path.join(root, file_name)
            file_list.append(file_path)
    return file_list


def read_csv_data_to_db(csv_path: str, con, db_name: str, symbols: Union[List, None] = None):
    """
    读取数据
    :param csv_path: csv文件路径
    :param con: 数据库对象
    :param db_name: 数据库标记
    :param symbols: 指定入库的合约，若为None, 所有合约都导入数据库
    :return:
    """
    data_name: Dict = {
        'day_bar':  {'mysql': to_save_mysql_day, 'dolphindb': to_save_dolphindb_day},
        'min_bar': {'mysql': to_save_mysql_bar, 'dolphindb': to_save_dolphindb_bar},
        'tick_data': {'mysql': to_save_mysql_tick, 'dolphindb': to_save_dolphindb_tick}
    }
    ExchangeId = csv_path.split('\\')[-1]  # 交易所代码

    file_list = get_all_files(csv_path)
    print(f'文件列表：{file_list}')
    for rar_file in file_list:
        if rar_file.endswith('rar'):
            save_path = os.path.dirname(rar_file)
            with rarfile.RarFile(rar_file) as rf:
                rf.extractall(save_path)
                print(f'{rar_file} ==> 解压完成')

            for name in ['day_bar', 'min_bar', 'tick_data']:
                csv_path = os.path.join(rar_file[:-4], name)
                csv_list = get_all_files(csv_path)
                for csv_file in csv_list:
                    if symbols is None:
                        df = pd.read_csv(csv_file, encoding='utf-8')
                        df['ExchangeID'] = ExchangeId
                        # 入库
                        data_name[name][db_name](con, df)
                    else:
                        file_name = csv_file.split('\\')[-1].split('.')[0]
                        if file_name in symbols:
                            df = pd.read_csv(csv_file, encoding='utf-8')
                            df['ExchangeID'] = ExchangeId
                            # 入库
                            data_name[name][db_name](con, df)
            remove_directory(rar_file.split('.')[0])
        else:
            print(f'{rar_file} ==> 文件格式不正确')


def to_mysql(csv_path: str, host: str, port: int, userid: str, password: str, symbols: Union[List, None] = None,
             database: str = 'lhxt'):
    """
    mysql数据库
    :param csv_path: csv路径
    :param host: IP地址
    :param port: 端口
    :param userid: 用户名
    :param password: 密码
    :param symbols: 指定入库的合约，若为None, 所有合约都导入数据库（所有合约：小写字母+数字）
    :param database: 数据库
    :return:
    """
    dbName = 'mysql'  # 数据库
    s = connect_mysql(userid, password, host, port, database)
    read_csv_data_to_db(csv_path, s, dbName, symbols)


def to_dolphindb(csv_path: str, host: str, port: int, userid: str, password: str, symbols: Union[List, None] = None):
    """
    dolphindb数据库
    :param csv_path: csv路径
    :param host: IP地址
    :param port: 端口
    :param userid: 用户名
    :param password: 密码
    :param symbols: 指定入库的合约，若为None, 所有合约都导入数据库（合约：小写字母+数字）
    :return:
    """
    dbName = 'dolphindb'  # 数据库
    s = ddb.session()
    s.connect(host, port, userid, password)
    read_csv_data_to_db(csv_path, s, dbName, symbols)
    s.close()


if __name__ == '__main__':
    rarfile.UNRAR_TOOL = r'C:\Users\jiaoyulong\Downloads\UnRAR.exe'  # 设置UnRAR.exe路径
    csvPath = r'E:\Future_data\CFFEX'  # 导入CFFEX数据（文件路径末尾必须是交易所代码）

    symbols = ['sa2501', 'ur2501']  # csvPath路径下的指定合约入库

    # 1.dolphindb
    # to_dolphindb(csvPath, 'localhost', 8848, 'admin', '123456')  # 所有合约入库
    # to_dolphindb(csvPath, 'localhost', 8848, 'admin', '123456', symbols)  # 指定合约入库

    # 2.mysql
    to_mysql(csvPath, 'localhost', 3306, 'root', '123456')
    # to_mysql(csvPath, 'localhost', 3306, 'root', '123456', symbols)




